package com.lan.ai;

import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
import com.lan.ai.mutimodelEmbedding.MutimodeMilvusVectorStore;
import io.micrometer.observation.ObservationRegistry;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusServiceClientConnectionDetails;
import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusServiceClientProperties;
import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreAutoConfiguration;
import org.springframework.ai.autoconfigure.vectorstore.milvus.MilvusVectorStoreProperties;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.vectorstore.MilvusVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.StringUtils;

import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @author lanyanhua
 * @date 2024/12/24 11:47
 * @description 需要排除 MilvusVectorStoreAutoConfiguration.class 指定OpenAiEmbeddingModel embeddingModel嵌入模型
 */
@Slf4j
@Configuration
@EnableAutoConfiguration(exclude = {MilvusVectorStoreAutoConfiguration.class})
@ConditionalOnClass({MilvusVectorStore.class, EmbeddingModel.class})
@EnableConfigurationProperties({MilvusServiceClientProperties.class, MilvusVectorStoreProperties.class})
public class VectorStoreConfig {

    @Bean
    PropertiesMilvusServiceClientConnectionDetails milvusServiceClientConnectionDetails(
            MilvusServiceClientProperties properties) {
        return new PropertiesMilvusServiceClientConnectionDetails(properties);
    }

    @Bean
    public MutimodeMilvusVectorStore imgVectorStore(MilvusServiceClient imgMilvusClient, DashScopeEmbeddingModel dashScopeMultiModalEmbedding,
                                                    MilvusVectorStoreProperties properties,
                                                    ObjectProvider<ObservationRegistry> observationRegistry,
                                                    ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {
        //todo 分割 保证每次都是一个，不然太长了模型不支持
        BatchingStrategy batchingStrategy = (documents) -> documents.stream().map(List::of).collect(Collectors.toList());

        MutimodeMilvusVectorStore.MilvusVectorStoreConfig config = new MutimodeMilvusVectorStore.MilvusVectorStoreConfig(
                "img_db", "img_coll",
                properties.getEmbeddingDimension(),
                IndexType.valueOf(properties.getIndexType().name()),
                MetricType.valueOf(properties.getMetricType().name()),
                properties.getIndexParameters()
        );

        return new MutimodeMilvusVectorStore(imgMilvusClient, dashScopeMultiModalEmbedding, config, properties.isInitializeSchema(),
                batchingStrategy, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
                customObservationConvention.getIfAvailable(() -> null));
    }


    @Bean
    public MilvusServiceClient imgMilvusClient(MilvusServiceClientProperties clientProperties
            , MilvusServiceClientConnectionDetails connectionDetails) {

        var builder = ConnectParam.newBuilder()
                .withHost(connectionDetails.getHost())
                .withPort(connectionDetails.getPort())
                .withDatabaseName("img_db")
                .withConnectTimeout(clientProperties.getConnectTimeoutMs(), TimeUnit.MILLISECONDS)
                .withKeepAliveTime(clientProperties.getKeepAliveTimeMs(), TimeUnit.MILLISECONDS)
                .withKeepAliveTimeout(clientProperties.getKeepAliveTimeoutMs(), TimeUnit.MILLISECONDS)
                .withRpcDeadline(clientProperties.getRpcDeadlineMs(), TimeUnit.MILLISECONDS)
                .withSecure(clientProperties.isSecure())
                .withIdleTimeout(clientProperties.getIdleTimeoutMs(), TimeUnit.MILLISECONDS)
                .withAuthorization(clientProperties.getUsername(), clientProperties.getPassword());

        return getMilvusServiceClient(clientProperties, builder);
    }

    @NotNull
    private MilvusServiceClient getMilvusServiceClient(MilvusServiceClientProperties clientProperties, ConnectParam.Builder builder) {
        if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getUri())) {
            builder.withUri(clientProperties.getUri());
        }

        if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getToken())) {
            builder.withToken(clientProperties.getToken());
        }

        if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getClientKeyPath())) {
            builder.withClientKeyPath(clientProperties.getClientKeyPath());
        }

        if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getClientPemPath())) {
            builder.withClientPemPath(clientProperties.getClientPemPath());
        }

        if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getCaPemPath())) {
            builder.withCaPemPath(clientProperties.getCaPemPath());
        }

        if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getServerPemPath())) {
            builder.withServerPemPath(clientProperties.getServerPemPath());
        }

        if (clientProperties.isSecure() && StringUtils.hasText(clientProperties.getServerName())) {
            builder.withServerName(clientProperties.getServerName());
        }

        return new MilvusServiceClient(builder.build());
    }



    @Bean
    BatchingStrategy milvusBatchingStrategy() {
        return new TokenCountBatchingStrategy();
    }
    @Bean
    public MilvusVectorStore vectorStore(MilvusServiceClient milvusClient, DashScopeEmbeddingModel embeddingModel,
                                         MilvusVectorStoreProperties properties, BatchingStrategy batchingStrategy,
                                         ObjectProvider<ObservationRegistry> observationRegistry,
                                         ObjectProvider<VectorStoreObservationConvention> customObservationConvention) {

        MilvusVectorStore.MilvusVectorStoreConfig config = MilvusVectorStore.MilvusVectorStoreConfig.builder()
                .withCollectionName(properties.getCollectionName())
                .withDatabaseName(properties.getDatabaseName())
                .withIndexType(IndexType.valueOf(properties.getIndexType().name()))
                .withMetricType(MetricType.valueOf(properties.getMetricType().name()))
                .withIndexParameters(properties.getIndexParameters())
                .withEmbeddingDimension(properties.getEmbeddingDimension())
                .build();

        return new MilvusVectorStore(milvusClient, embeddingModel, config, properties.isInitializeSchema(),
                batchingStrategy, observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP),
                customObservationConvention.getIfAvailable(() -> null));
    }


    @Bean
    public MilvusServiceClient milvusClient(MilvusVectorStoreProperties serverProperties,
                                            MilvusServiceClientProperties clientProperties, MilvusServiceClientConnectionDetails connectionDetails) {

        var builder = ConnectParam.newBuilder()
                .withHost(connectionDetails.getHost())
                .withPort(connectionDetails.getPort())
                .withDatabaseName(serverProperties.getDatabaseName())
                .withConnectTimeout(clientProperties.getConnectTimeoutMs(), TimeUnit.MILLISECONDS)
                .withKeepAliveTime(clientProperties.getKeepAliveTimeMs(), TimeUnit.MILLISECONDS)
                .withKeepAliveTimeout(clientProperties.getKeepAliveTimeoutMs(), TimeUnit.MILLISECONDS)
                .withRpcDeadline(clientProperties.getRpcDeadlineMs(), TimeUnit.MILLISECONDS)
                .withSecure(clientProperties.isSecure())
                .withIdleTimeout(clientProperties.getIdleTimeoutMs(), TimeUnit.MILLISECONDS)
                .withAuthorization(clientProperties.getUsername(), clientProperties.getPassword());

        return getMilvusServiceClient(clientProperties, builder);
    }

    static class PropertiesMilvusServiceClientConnectionDetails implements MilvusServiceClientConnectionDetails {

        private final MilvusServiceClientProperties properties;

        PropertiesMilvusServiceClientConnectionDetails(MilvusServiceClientProperties properties) {
            this.properties = properties;
        }

        @Override
        public String getHost() {
            return this.properties.getHost();
        }

        @Override
        public int getPort() {
            return this.properties.getPort();
        }

    }

}
